package org.kirin.ocr.service;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.nio.charset.Charset;
import java.text.SimpleDateFormat;
import java.util.List;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: Kirin
 * @Date: 2024/12/3 00:23
 * @Description:
 **/
@Slf4j
public class MaJiongOcrService {
    public void ocr(String path) {
        ITesseract tesseract = new Tesseract();
        tesseract.setDatapath("src\\main\\resources\\tess4j");
        tesseract.setLanguage("chi_sim+eng");
//        tesseract.setTessVariable("tessedit_char_whitelist", "1234567890");

        Map<String, List<Rectangle>> map = new HashMap<>();


        List<Rectangle> rectangles = this.getRectangles();
        AtomicInteger i = new AtomicInteger();
        FileUtil.loopFiles(path).forEach(file -> {
            rectangles.forEach(rectangle -> {
                try {
                    BufferedImage originalImage = ImageIO.read(file);
                    BufferedImage splitImage = generateImage(originalImage, rectangle);
                    splitImage = this.binarize(splitImage, 128);

                    // 将截取后的图像保存到文件
                    File outputFile = new File("D:\\zmq\\code\\ocr\\src\\main\\resources\\tess4j\\image\\test\\" + i.getAndIncrement() + ".jpg"); // 替换为你想要保存的路径
                    ImageIO.write(splitImage, "jpg", outputFile);

//                    System.out.println(tesseract.doOCR(splitImage, rectangle));
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            });
        });
    }

    public List<Rectangle> getRectangles() {
        File dirPath = new File("src\\main\\resources\\tess4j\\majiang\\");

        List<Rectangle> rectangles = new ArrayList<>();
        FileUtil.loopFiles(dirPath)
                .forEach(file -> {
                    List<String> lines = FileUtil.readLines(file, Charset.defaultCharset());
                    for (String line : lines) {
                        List<String> points = StrUtil.split(line, ",");
                        int x1 = Integer.parseInt(points.get(0));
                        int y1 = Integer.parseInt(points.get(1));
                        int x2 = Integer.parseInt(points.get(2));
                        int y2 = Integer.parseInt(points.get(3));
                        Rectangle rectangle = new Rectangle(x1, y1, x2 - x1, y2 - y1);
                        rectangles.add(rectangle);
                    }
                });
        return rectangles;
    }

    /**
     * 截取图像
     * @param originalImage
     * @param rectangle
     */
    private BufferedImage generateImage(BufferedImage originalImage, Rectangle rectangle) {
        // 创建一个新的BufferedImage，大小为指定的矩形范围
        BufferedImage croppedImage = new BufferedImage(rectangle.width, rectangle.height, BufferedImage.TYPE_INT_RGB);

        // 获取新图像的图形上下文
        Graphics2D g2d = croppedImage.createGraphics();

        // 将原始图像中指定矩形范围内的内容绘制到新图像上
        g2d.drawImage(originalImage, 0, 0, rectangle.width, rectangle.height, rectangle.x, rectangle.y,
                rectangle.x + rectangle.width, rectangle.y + rectangle.height, null);

        // 释放图形上下文资源
        g2d.dispose();

        return croppedImage;
    }

    /**
     * 二值化处理
     * @param image
     * @param threshold
     * @return
     */
    public BufferedImage binarize(BufferedImage image, int threshold) {
        int width = image.getWidth();
        int height = image.getHeight();
        BufferedImage binaryImage = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_BINARY);
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int argb = image.getRGB(x, y);
                int alpha = (argb >> 24) & 0xff;
                int red = (argb >> 16) & 0xff;
                int green = (argb >> 8) & 0xff;
                int blue = argb & 0xff;
                int gray = (int)(0.299 * red + 0.587 * green + 0.114 * blue);
                int newPixel = gray > threshold? 0xFFFFFFFF : 0xFF000000;
                binaryImage.setRGB(x, y, newPixel);
            }
        }
        return binaryImage;
    }

    /**
     * 画出识别区
     * @param path
     * @throws Exception
     */
    public void draw(String path) throws Exception {
        File imageFile = new File(path);
        BufferedImage bi = ImageIO.read(imageFile);
        // 复制原始图像用于绘制框框
        BufferedImage outputImage = new BufferedImage(bi.getWidth(), bi.getHeight(), bi.getType());
        Graphics2D g2 = outputImage.createGraphics();
        g2.drawImage(bi, 0, 0, null);
        g2.dispose();

        // 绘制识别区域的框框并保存输出图像
        Graphics2D g2d = outputImage.createGraphics();
        g2d.setColor(Color.RED);
        g2d.setStroke(new BasicStroke(2));

        this.getRectangles().forEach(rectangle -> {
            g2d.drawRect(rectangle.x, rectangle.y, rectangle.width, rectangle.height);
        });

        g2d.dispose();

        // 在识别的同级目录下创建输出文件，文件名格式为：源文件名称_年月日时分秒
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMddHHmmss");
        String outputFileName = imageFile.getName().substring(0, imageFile.getName().lastIndexOf('.')) + "_" + sdf.format(new Date()) + ".png";
        File outputFile = new File(imageFile.getParent(), outputFileName);
        ImageIO.write(outputImage, "png", outputFile);
    }

    public static void main(String[] args) throws Exception {
        MaJiongOcrService service = new MaJiongOcrService();
        service.ocr("D:\\zmq\\code\\ocr\\src\\main\\resources\\tess4j\\image\\test");
//        service.draw("D:\\zmq\\code\\ocr\\src\\main\\resources\\tess4j\\image\\test\\test-1.jpg");
    }
}
